#include "PredicatePushdownOptimizer.h"
#include "AST.h"
#include <iostream>
#include <algorithm>

namespace MiniDB {

PredicatePushdownOptimizer::PredicatePushdownOptimizer(SymbolTable& symbol_table)
    : symbol_table_(symbol_table) {
}

std::unique_ptr<ExecutionPlan> PredicatePushdownOptimizer::optimize(std::unique_ptr<ExecutionPlan> plan) {
    if (!plan || !plan->getRoot()) {
        return plan;
    }

    std::cout << "开始应用谓词下推优化..." << std::endl;

    // 优化根算子
    auto optimized_root = optimizeOperator(plan->getRoot());

    // 创建新的执行计划
    auto optimized_plan = std::make_unique<ExecutionPlan>();
    optimized_plan->setRoot(optimized_root);

    std::cout << "谓词下推优化完成" << std::endl;
    return optimized_plan;
}

std::shared_ptr<PlanOperator> PredicatePushdownOptimizer::optimizeOperator(std::shared_ptr<PlanOperator> op) {
    if (!op) {
        return nullptr;
    }

    // 根据算子类型进行优化
    if (op->getType() == "Filter") {
        return optimizeFilter(std::static_pointer_cast<FilterOperator>(op));
    } else if (op->getType() == "Project") {
        return optimizeProject(std::static_pointer_cast<ProjectOperator>(op));
    } else if (op->getType() == "SeqScan") {
        // 扫描算子已经是叶子节点，无需进一步优化
        return op;
    } else {
        // 对于其他类型的算子，递归优化子算子
        if (auto child = op->getChild()) {
            auto optimized_child = optimizeOperator(child);
            // 这里需要根据具体的算子类型更新子算子
            // 简化起见，先返回原算子
            return op;
        }
        return op;
    }
}

std::shared_ptr<PlanOperator> PredicatePushdownOptimizer::optimizeFilter(std::shared_ptr<FilterOperator> filter_op) {
    auto child = filter_op->getChild();
    auto condition = filter_op->getCondition();

    if (!child || !condition) {
        return filter_op;
    }

    // 递归优化子算子
    auto optimized_child = optimizeOperator(child);

    // 检查子算子是否是扫描算子
    if (optimized_child->getType() == "SeqScan") {
        auto seq_scan = std::static_pointer_cast<SeqScanOperator>(optimized_child);
        const std::string& table_name = seq_scan->getTableName();

        // 检查过滤条件是否可以完全下推到扫描算子
        if (canPushDownToTable(condition, table_name)) {
            std::cout << "将过滤条件完全下推到表 " << table_name << " 的扫描算子" << std::endl;
            // 创建带有过滤条件的扫描算子
            auto filtered_scan = std::make_shared<SeqScanOperator>(table_name);
            return createFilter(std::static_pointer_cast<PlanOperator>(filtered_scan), condition);
        } else {
            // 尝试提取可以下推的部分
            auto pushable_predicates = extractPushablePredicates(condition, table_name);
            if (!pushable_predicates.empty()) {
                std::cout << "将部分过滤条件下推到表 " << table_name << " 的扫描算子" << std::endl;

                // 创建带有部分过滤条件的扫描算子
                auto filtered_scan = std::make_shared<SeqScanOperator>(table_name);
                auto pushable_condition = combinePredicates(pushable_predicates);
                auto scan_with_filter = createFilter(std::static_pointer_cast<PlanOperator>(filtered_scan), pushable_condition);

                // 剩余条件保持在原Filter算子中
                return createFilter(scan_with_filter, condition);
            }
        }
    }

    // 如果无法下推，返回优化后的算子
    if (optimized_child != child) {
        return createFilter(optimized_child, condition);
    }

    return filter_op;
}

std::shared_ptr<PlanOperator> PredicatePushdownOptimizer::optimizeProject(std::shared_ptr<ProjectOperator> project_op) {
    auto child = project_op->getChild();

    if (!child) {
        return project_op;
    }

    // 递归优化子算子
    auto optimized_child = optimizeOperator(child);

    if (optimized_child != child) {
        // 创建新的Project算子
        return std::make_shared<ProjectOperator>(optimized_child, project_op->getProjections());
    }

    return project_op;
}

bool PredicatePushdownOptimizer::canPushDownToTable(std::shared_ptr<Expression> expr, const std::string& table_name) {
    return involvesOnlyTableColumns(expr, table_name);
}

std::vector<std::shared_ptr<Expression>> PredicatePushdownOptimizer::extractPushablePredicates(
    std::shared_ptr<Expression> expr, const std::string& table_name) {

    std::vector<std::shared_ptr<Expression>> pushable_predicates;

    // 如果是AND表达式，分解为多个条件
    if (auto binary_expr = std::dynamic_pointer_cast<BinaryExpression>(expr)) {
        if (binary_expr->operator_ == "AND") {
            // 递归处理左右子表达式
            auto left_pushable = extractPushablePredicates(binary_expr->left, table_name);
            auto right_pushable = extractPushablePredicates(binary_expr->right, table_name);

            // 合并可下推的谓词
            pushable_predicates.insert(pushable_predicates.end(),
                                     left_pushable.begin(), left_pushable.end());
            pushable_predicates.insert(pushable_predicates.end(),
                                     right_pushable.begin(), right_pushable.end());
        } else {
            // 对于其他二元表达式，检查整体是否可以下推
            if (canPushDownToTable(expr, table_name)) {
                pushable_predicates.push_back(expr);
            }
        }
    } else {
        // 对于非AND表达式，检查是否可以下推
        if (canPushDownToTable(expr, table_name)) {
            pushable_predicates.push_back(expr);
        }
    }

    return pushable_predicates;
}

std::shared_ptr<Expression> PredicatePushdownOptimizer::extractTableSpecificPredicate(
    std::shared_ptr<Expression> expr, const std::string& table_name) {

    // 如果表达式只涉及指定表的列，返回该表达式
    if (involvesOnlyTableColumns(expr, table_name)) {
        return expr;
    }

    // 如果是AND表达式，尝试提取涉及指定表的子表达式
    if (auto binary_expr = std::dynamic_pointer_cast<BinaryExpression>(expr)) {
        if (binary_expr->operator_ == "AND") {
            auto left_specific = extractTableSpecificPredicate(binary_expr->left, table_name);
            auto right_specific = extractTableSpecificPredicate(binary_expr->right, table_name);

            std::vector<std::shared_ptr<Expression>> parts;
            if (left_specific) parts.push_back(left_specific);
            if (right_specific) parts.push_back(right_specific);

            if (parts.size() == 1) {
                return parts[0];
            } else if (parts.size() == 2) {
                return std::make_shared<BinaryExpression>(parts[0], "AND", parts[1]);
            }
        }
    }

    return nullptr;
}

std::shared_ptr<Expression> PredicatePushdownOptimizer::combinePredicates(
    const std::vector<std::shared_ptr<Expression>>& predicates) {

    if (predicates.empty()) {
        return nullptr;
    }

    if (predicates.size() == 1) {
        return predicates[0];
    }

    // 使用AND连接多个谓词
    auto result = predicates[0];
    for (size_t i = 1; i < predicates.size(); ++i) {
        result = std::make_shared<BinaryExpression>(result, "AND", predicates[i]);
    }

    return result;
}

bool PredicatePushdownOptimizer::involvesOnlyTableColumns(std::shared_ptr<Expression> expr, const std::string& table_name) {
    std::unordered_set<std::string> column_refs;
    return checkColumnReferences(expr, table_name, column_refs);
}

bool PredicatePushdownOptimizer::checkColumnReferences(std::shared_ptr<Expression> expr,
                                                      const std::string& table_name,
                                                      std::unordered_set<std::string>& column_refs) {
    if (auto column_ref = std::dynamic_pointer_cast<ColumnReference>(expr)) {
        // 检查列引用是否属于指定表
        if (!column_ref->table_name.empty() && column_ref->table_name != table_name) {
            return false; // 引用了其他表的列
        }
        column_refs.insert(column_ref->column_name);
        return true;
    }

    if (auto binary_expr = std::dynamic_pointer_cast<BinaryExpression>(expr)) {
        return checkColumnReferences(binary_expr->left, table_name, column_refs) &&
               checkColumnReferences(binary_expr->right, table_name, column_refs);
    }

    if (auto unary_expr = std::dynamic_pointer_cast<UnaryExpression>(expr)) {
        return checkColumnReferences(unary_expr->operand, table_name, column_refs);
    }

    if (auto func_call = std::dynamic_pointer_cast<FunctionCall>(expr)) {
        for (auto& arg : func_call->arguments) {
            if (!checkColumnReferences(arg, table_name, column_refs)) {
                return false;
            }
        }
        return true;
    }

    // 对于字面量等其他表达式，返回true（不涉及列引用）
    return true;
}

std::unordered_set<std::string> PredicatePushdownOptimizer::getReferencedTables(std::shared_ptr<Expression> expr) {
    std::unordered_set<std::string> tables;

    if (auto column_ref = std::dynamic_pointer_cast<ColumnReference>(expr)) {
        if (!column_ref->table_name.empty()) {
            tables.insert(column_ref->table_name);
        }
    }

    if (auto binary_expr = std::dynamic_pointer_cast<BinaryExpression>(expr)) {
        auto left_tables = getReferencedTables(binary_expr->left);
        auto right_tables = getReferencedTables(binary_expr->right);
        tables.insert(left_tables.begin(), left_tables.end());
        tables.insert(right_tables.begin(), right_tables.end());
    }

    if (auto unary_expr = std::dynamic_pointer_cast<UnaryExpression>(expr)) {
        auto operand_tables = getReferencedTables(unary_expr->operand);
        tables.insert(operand_tables.begin(), operand_tables.end());
    }

    if (auto func_call = std::dynamic_pointer_cast<FunctionCall>(expr)) {
        for (auto& arg : func_call->arguments) {
            auto arg_tables = getReferencedTables(arg);
            tables.insert(arg_tables.begin(), arg_tables.end());
        }
    }

    return tables;
}

std::shared_ptr<FilterOperator> PredicatePushdownOptimizer::createFilter(std::shared_ptr<PlanOperator> child,
                                                                        std::shared_ptr<Expression> condition) {
    return std::make_shared<FilterOperator>(child, condition);
}

std::shared_ptr<PlanOperator> PredicatePushdownOptimizer::removeRedundantFilters(std::shared_ptr<PlanOperator> op) {
    // 简化实现：移除连续的相同过滤条件
    // 在实际实现中，这里应该有更复杂的逻辑来检测和移除冗余过滤器

    if (op->getType() == "Filter") {
        auto filter_op = std::static_pointer_cast<FilterOperator>(op);
        auto child = filter_op->getChild();

        if (child && child->getType() == "Filter") {
            auto child_filter = std::static_pointer_cast<FilterOperator>(child);
            // 如果两个过滤条件相同，可以合并
            // 这里简化处理，直接返回子过滤器
            return child;
        }
    }

    return op;
}

} // namespace MiniDB
